// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <tests/compute/kernels/test_util.h>

#include <cstdint>
#include <memory>
#include <string>

#include <nebula/core/array.h>
#include <nebula/array/validate.h>
#include <nebula/core/chunked_array.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/exec.h>
#include <nebula/compute/function.h>
#include <nebula/compute/registry.h>
#include <nebula/core/datum.h>

#include <nebula/core/table.h>
#include <nebula/testing/fixed_width_test_util.h>
#include <nebula/testing/ktest_util.h>

namespace nebula {
    namespace compute {

        namespace {

            template<typename T>
            DatumVector GetDatums(const std::vector<T> &inputs) {
                DatumVector datums;
                for (const auto &input: inputs) {
                    datums.emplace_back(input);
                }
                return datums;
            }

            template<typename... SliceArgs>
            DatumVector SliceArrays(const DatumVector &inputs, SliceArgs... slice_args) {
                DatumVector sliced;
                for (const auto &input: inputs) {
                    if (input.is_array()) {
                        sliced.push_back(*input.make_array()->slice(slice_args...));
                    } else {
                        sliced.push_back(input);
                    }
                }
                return sliced;
            }

            ScalarVector GetScalars(const DatumVector &inputs, int64_t index) {
                ScalarVector scalars;
                for (const auto &input: inputs) {
                    if (input.is_array()) {
                        scalars.push_back(*input.make_array()->get_scalar(index));
                    } else {
                        scalars.push_back(input.scalar());
                    }
                }
                return scalars;
            }

        }  // namespace

        void check_scalar_non_recursive(const std::string &func_name, const DatumVector &inputs,
                                        const Datum &expected, const FunctionOptions *options) {
            ASSERT_OK_AND_ASSIGN(Datum out, call_function(func_name, inputs, options));
            ValidateOutput(out);
            AssertDatumsEqual(expected, out, /*verbose=*/true);
        }

        void check_scalar(std::string func_name, const ScalarVector &inputs,
                          std::shared_ptr<Scalar> expected, const FunctionOptions *options) {
            ASSERT_OK_AND_ASSIGN(Datum out, call_function(func_name, GetDatums(inputs), options));
            ValidateOutput(out);
            if (!out.scalar()->equals(*expected)) {
                std::string summary = func_name + "(";
                for (const auto &input: inputs) {
                    summary += input->to_string() + ",";
                }
                summary.back() = ')';

                summary += " = " + out.scalar()->to_string() + " != " + expected->to_string();

                if (!out.type()->equals(expected->type)) {
                    summary += " (types differed: " + out.type()->to_string() + " vs " +
                               expected->type->to_string() + ")";
                }

                FAIL() << summary;
            }
        }

        void check_scalar(std::string func_name, const DatumVector &inputs, Datum expected_datum,
                          const FunctionOptions *options) {
            check_scalar_non_recursive(func_name, inputs, expected_datum, options);

            if (expected_datum.is_scalar()) return;
            ASSERT_TRUE(expected_datum.is_array())
                                        << "check_scalar is only implemented for scalar/array expected values";
            auto expected = expected_datum.make_array();

            // check for at least 1 array, and make sure the others are of equal length
            bool has_array = false;
            for (const auto &input: inputs) {
                if (input.is_array()) {
                    ASSERT_EQ(input.array()->length, expected->length());
                    has_array = true;
                }
            }
            ASSERT_TRUE(has_array) << "Must have at least 1 array input to have an array output";

            // Check all the input scalars
            for (int64_t i = 0; i < expected->length(); ++i) {
                check_scalar(func_name, GetScalars(inputs, i), *expected->get_scalar(i), options);
            }

            // Since it's a scalar function, calling it on sliced inputs should
            // result in the sliced expected output.
            const auto slice_length = expected->length() / 3;
            if (slice_length > 0) {
                check_scalar_non_recursive(func_name, SliceArrays(inputs, 0, slice_length),
                                           expected->slice(0, slice_length), options);

                check_scalar_non_recursive(func_name, SliceArrays(inputs, slice_length, slice_length),
                                           expected->slice(slice_length, slice_length), options);

                check_scalar_non_recursive(func_name, SliceArrays(inputs, 2 * slice_length),
                                           expected->slice(2 * slice_length), options);
            }

            // Should also work with an empty slice
            check_scalar_non_recursive(func_name, SliceArrays(inputs, 0, 0), expected->slice(0, 0),
                                       options);

            // Ditto with ChunkedArray inputs
            if (slice_length > 0) {
                DatumVector chunked_inputs;
                chunked_inputs.reserve(inputs.size());
                for (const auto &input: inputs) {
                    if (input.is_array()) {
                        auto ar = input.make_array();
                        auto ar_chunked = std::make_shared<ChunkedArray>(
                                ArrayVector{ar->slice(0, slice_length), ar->slice(slice_length)});
                        chunked_inputs.push_back(ar_chunked);
                    } else {
                        chunked_inputs.push_back(input.scalar());
                    }
                }
                ArrayVector expected_chunks{expected->slice(0, slice_length),
                                            expected->slice(slice_length)};

                ASSERT_OK_AND_ASSIGN(Datum out,
                                     call_function(func_name, GetDatums(chunked_inputs), options));
                ValidateOutput(out);
                auto chunked = out.chunked_array();
                (void) chunked;
                AssertDatumsEqual(std::make_shared<ChunkedArray>(expected_chunks), out);
            }
        }

        Datum CheckDictionaryNonRecursive(const std::string &func_name, const DatumVector &args,
                                          bool result_is_encoded) {
            EXPECT_OK_AND_ASSIGN(Datum actual, call_function(func_name, args));
            ValidateOutput(actual);

            DatumVector decoded_args;
            decoded_args.reserve(args.size());
            for (const auto &arg: args) {
                if (arg.type()->id() == Type::DICTIONARY) {
                    const auto &to_type = turbo::checked_cast<const DictionaryType &>(*arg.type()).get_value_type();
                    EXPECT_OK_AND_ASSIGN(auto decoded, Cast(arg, to_type));
                    decoded_args.push_back(decoded);
                } else {
                    decoded_args.push_back(arg);
                }
            }
            EXPECT_OK_AND_ASSIGN(Datum expected, call_function(func_name, decoded_args));

            if (result_is_encoded) {
                EXPECT_EQ(Type::DICTIONARY, actual.type()->id())
                                    << "turbo::Result should have been dictionary-encoded";
                // Decode before comparison - we care about equivalent not identical results
                const auto &to_type =
                        turbo::checked_cast<const DictionaryType &>(*actual.type()).get_value_type();
                EXPECT_OK_AND_ASSIGN(auto decoded, Cast(actual, to_type));
                AssertDatumsApproxEqual(expected, decoded, /*verbose=*/true);
            } else {
                AssertDatumsApproxEqual(expected, actual, /*verbose=*/true);
            }
            return actual;
        }

        void CheckDictionary(const std::string &func_name, const DatumVector &args,
                             bool result_is_encoded) {
            auto actual = CheckDictionaryNonRecursive(func_name, args, result_is_encoded);

            if (actual.is_scalar()) return;
            ASSERT_TRUE(actual.is_array());
            ASSERT_GE(actual.length(), 0);

            // Check all scalars
            for (int64_t i = 0; i < actual.length(); i++) {
                CheckDictionaryNonRecursive(func_name, GetDatums(GetScalars(args, i)),
                                            result_is_encoded);
            }

            // Check slices of the input
            const auto slice_length = actual.length() / 3;
            if (slice_length > 0) {
                CheckDictionaryNonRecursive(func_name, SliceArrays(args, 0, slice_length),
                                            result_is_encoded);
                CheckDictionaryNonRecursive(func_name, SliceArrays(args, slice_length, slice_length),
                                            result_is_encoded);
                CheckDictionaryNonRecursive(func_name, SliceArrays(args, 2 * slice_length),
                                            result_is_encoded);
            }

            // Check empty slice
            CheckDictionaryNonRecursive(func_name, SliceArrays(args, 0, 0), result_is_encoded);

            // Check chunked arrays
            if (slice_length > 0) {
                DatumVector chunked_args;
                chunked_args.reserve(args.size());
                for (const auto &arg: args) {
                    if (arg.is_array()) {
                        auto arr = arg.make_array();
                        ArrayVector chunks{arr->slice(0, slice_length), arr->slice(slice_length)};
                        chunked_args.push_back(std::make_shared<ChunkedArray>(std::move(chunks)));
                    } else {
                        chunked_args.push_back(arg);
                    }
                }
                CheckDictionaryNonRecursive(func_name, chunked_args, result_is_encoded);
            }
        }

        void CheckScalarUnary(std::string func_name, Datum input, Datum expected,
                              const FunctionOptions *options) {
            DatumVector input_vector = {std::move(input)};
            check_scalar(std::move(func_name), input_vector, expected, options);
        }

        void CheckScalarUnary(std::string func_name, std::shared_ptr<DataType> in_ty,
                              std::string json_input, std::shared_ptr<DataType> out_ty,
                              std::string json_expected, const FunctionOptions *options) {
            CheckScalarUnary(std::move(func_name), assert_array_from_json(in_ty, json_input),
                             assert_array_from_json(out_ty, json_expected), options);
        }

        void CheckVectorUnary(std::string func_name, Datum input, Datum expected,
                              const FunctionOptions *options) {
            ASSERT_OK_AND_ASSIGN(Datum actual, call_function(func_name, {input}, options));
            ValidateOutput(actual);
            AssertDatumsEqual(expected, actual, /*verbose=*/true);
        }

        void check_scalar_binary(std::string func_name, Datum left_input, Datum right_input,
                                 Datum expected, const FunctionOptions *options) {
            check_scalar(std::move(func_name), {left_input, right_input}, expected, options);
        }

        void CheckScalarBinaryCommutative(std::string func_name, Datum left_input,
                                          Datum right_input, Datum expected,
                                          const FunctionOptions *options) {
            check_scalar(func_name, {left_input, right_input}, expected, options);
            check_scalar(func_name, {right_input, left_input}, expected, options);
        }

        namespace {

            void ValidateOutputImpl(const ArrayData &output) {
                ASSERT_OK(::nebula::internal::ValidateArrayFull(output));
                TestInitialized(output);
            }

            void ValidateOutputImpl(const ChunkedArray &output) {
                ASSERT_OK(output.validate_full());
                for (const auto &chunk: output.chunks()) {
                    TestInitialized(*chunk);
                }
            }

            void ValidateOutputImpl(const RecordBatch &output) {
                ASSERT_OK(output.validate_full());
                for (const auto &column: output.column_data()) {
                    TestInitialized(*column);
                }
            }

            void ValidateOutputImpl(const Table &output) {
                ASSERT_OK(output.validate_full());
                for (const auto &column: output.columns()) {
                    for (const auto &chunk: column->chunks()) {
                        TestInitialized(*chunk);
                    }
                }
            }

            void ValidateOutputImpl(const Scalar &output) { ASSERT_OK(output.validate_full()); }

        }  // namespace

        void ValidateOutput(const Datum &output) {
            switch (output.kind()) {
                case Datum::ARRAY:
                    ValidateOutputImpl(*output.array());
                    break;
                case Datum::CHUNKED_ARRAY:
                    ValidateOutputImpl(*output.chunked_array());
                    break;
                case Datum::RECORD_BATCH:
                    ValidateOutputImpl(*output.record_batch());
                    break;
                case Datum::TABLE:
                    ValidateOutputImpl(*output.table());
                    break;
                case Datum::SCALAR:
                    ValidateOutputImpl(*output.scalar());
                    break;
                default:
                    break;
            }
        }

        void CheckDispatchBest(std::string func_name, std::vector<TypeHolder> original_values,
                               std::vector<TypeHolder> expected_equivalent_values) {
            ASSERT_OK_AND_ASSIGN(auto function, get_function_registry()->get_function(func_name));

            auto values = original_values;
            ASSERT_OK_AND_ASSIGN(auto actual_kernel, function->dispatch_best(&values));

            ASSERT_OK_AND_ASSIGN(auto expected_kernel,
                                 function->dispatch_exact(expected_equivalent_values));

            EXPECT_EQ(actual_kernel, expected_kernel)
                                << "  dispatch_best" << TypeHolder::to_string(original_values) << " => "
                                << actual_kernel->signature->to_string() << "\n"
                                << "  dispatch_exact" << TypeHolder::to_string(expected_equivalent_values) << " => "
                                << expected_kernel->signature->to_string();
            EXPECT_EQ(values.size(), expected_equivalent_values.size());
            for (size_t i = 0; i < values.size(); i++) {
                AssertTypeEqual(*values[i], *expected_equivalent_values[i]);
            }
        }

        void CheckDispatchFails(std::string func_name, std::vector<TypeHolder> types) {
            ASSERT_OK_AND_ASSIGN(auto function, get_function_registry()->get_function(func_name));
            ASSERT_NOT_OK(function->dispatch_best(&types));
            ASSERT_NOT_OK(function->dispatch_exact(types));
        }

    }  // namespace compute
}  // namespace nebula
